from ultralytics import YOLO
import ultralytics
import torch
import numpy as np
import os
import cv2


if __name__ == '__main__':

	device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
	print(device)
	print(torch.cuda.is_available())

    # Load a model
	model = YOLO("yolov8n-seg.yaml")  # build a new model from YAML
	model = YOLO("yolov8n-seg.pt")  # load a pretrained model (recommended for training)
	model = YOLO("yolov8n-seg.yaml").load("yolov8n.pt")  # build from YAML and transfer weights

    # Train the model
	results = model.train(data="robot-seg.yaml", epochs=100, imgsz=640)

	import importlib
	try:
		importlib.import_module("shutil")  # 尝试导入模块
		print(f"shutil模块已安装")
	except ImportError:
		print(f"shutil模块未安装，正在安装...")
		import pip
		pip.main(['install', "shutilwhich"])  # 使用 pip 安装模块

	# copy best model to current dir
	import shutil
	shutil.copy("./runs/segment/train10/best.pt", "./yolov8n-seg-AGI.pt")